# coding:utf-8
# __user__ = hiicy redldw
# __time__ = 2019/10/8
# __file__ = torchscript_
# __desc__ =
# 能在c++环境里run

import torch
import torch.jit

# class MyCell(torch.nn.Module):
#     def __init__(self):
#         super(MyCell, self).__init__()
#         self.linear = torch.nn.Linear(4, 4)
#
#     def forward(self, x, h):
#         new_h = torch.tanh(self.linear(x) + h)
#         return new_h, new_h
# # capture the definition of your model
# my_cell = MyCell()
x, h = torch.rand(3, 4), torch.rand(3, 4)


# # 创建了torch.jit.ScriptModule的一个实例
# traced_cell = torch.jit.trace(my_cell,(x,h))
# print(traced_cell.code)
# traced_cell(x, h)

class MyDecisionGate(torch.nn.Module):
    def forward(self, x):
        if x.sum() > 0:
            return x
        else:
            return -x


class MyCell(torch.nn.Module):
    def __init__(self, dg):
        super(MyCell, self).__init__()
        self.dg = dg
        self.linear = torch.nn.Linear(4, 4)

    def forward(self, x, h):
        new_h = torch.tanh(self.dg(self.linear(x)) + h)
        return new_h, new_h


# a script compiler 去编译自己的source python code
scripted_gate = torch.jit.script(MyDecisionGate())


# my_cell = MyCell(scripted_gate)
# traced_cell = torch.jit.script(my_cell)
# print(traced_cell)

# 混合脚本编译和跟踪
class MyRNNLoop(torch.nn.Module):
    def __init__(self):
        super(MyRNNLoop, self).__init__()
        # torch.jit.script will inline the code
        # for a traced module, and tracing will inline the code for a scripted module.
        self.cell = torch.jit.trace(MyCell(scripted_gate), (x, h))

    def forward(self, xs):
        h, y = torch.zeros(3, 4), torch.zeros(3, 4)
        for i in range(xs.size(0)):
            y, h = self.cell(xs[i], h)
        return y, h


# rnn_loop = torch.jit.script(MyRNNLoop())
# print(rnn_loop.code)


class WrapRNN(torch.nn.Module):
    def __init__(self):
        super(WrapRNN, self).__init__()
        self.loop = torch.jit.script(MyRNNLoop())

    def forward(self, xs):
        y, h = self.loop(xs)
        return torch.relu(y)

traced = torch.jit.trace(WrapRNN(), (torch.rand(10, 3, 4)))
print(traced.code)
# Saving and Loading models
# This format includes code, parameters, attributes, and debug information
traced.save(r'E:\memory\models\wrapped_rnn.zip')
# loaded = torch.jit.load('wrapped_rnn.zip')
# print(loaded)
# print(loaded.code)

# 通过在ScriptModule的子类上使用torch.jit.script批注（对于函数）或torch.jit.script_method批注（对于方法）来实现。 拥有注释的函数的主体将直接转换为TorchScript
@torch.jit.script
def foo(x, y):
    if x.max() > y.max():
        r = x
    else:
        r = y
    return r
class MyModule(torch.jit.ScriptModule):
    def __init__(self, N, M):
        super(MyModule, self).__init__()
        self.weight = torch.nn.Parameter(torch.rand(N, M))

    @torch.jit.script_method
    def forward(self, input):
        return self.weight.mv(input)
